/*
 * @Author: LibraStalker 1299640057@qq.com
 * @Date: 2023-01-27 23:44:51
 * @FilePath: slice.ts
 * @Description: 自己实现的slice算子
 */
/**
 * @file slice
 * @description slice op，目前 case 很单一，暂时仅支持遇到的 case：slice final dim
 * @example x = [[1,2,3,4],[5,6,7,8]] axes=[1] starts=[2] ends=[3] => out [3,7]
 */

import { genGLSLArr, ArrTypeEnum, getValueFromArrByIndex } from '../../utils/dataProcess';


let current_index_num = 0; // 这里因为ESLint不能在函数中定义函数才用全局变量的
function getIndexArrFromShape(shape: Number[]) {
    const arr = [];
    if (shape.length === 1) {
        for (let i = 0; i < shape[0]; i++, current_index_num++) {
            arr.push(current_index_num);
        }
        return arr;
    }
    for (let i = 0; i < shape[0]; i++) {
        arr.push(getIndexArrFromShape(shape.slice(1, shape.length)));
    }
    return arr;
}

function getSliceIndex(
    {
        axes,
        starts,
        ends,
        decrease_axis
    },
    {
        _index_arr = undefined,
        shape = undefined
    }
) {
    let index_arr;
    if (_index_arr === undefined) {
        if (shape === undefined) { // 默认值为从0到总长度
            return;
        }
        current_index_num = 0;
        index_arr = getIndexArrFromShape(shape);
    } else {
        index_arr = _index_arr;
    }
    const slicers = [];
    // slicers.push([undefined, undefined]); // !!!!!!!测试用
    // decrease_axis.pop(); // !!!!!!!测试用
    const if_decrease_axis = [];
    for (let i = 0; i <= Math.max(...axes.concat(decrease_axis)); i++) { // 取axes和decrease_axis中所有元素的最大值，防止漏处理之后的维度
        if (axes.includes(i)) {
            const axis_index = axes.indexOf(i);
            slicers.push([starts[axis_index], ends[axis_index]]); // 当前维度切片的起止位置
        }
        else {
            slicers.push([undefined, undefined]); // 当前维度不需要切片，直接[...].slice(undefined, undefined)就是返回原数组的复制
        }
        if_decrease_axis.push(decrease_axis.includes(i + 1)); // 这里意味着合并i到i+1这两个维度
    }
    return getSliceIndexRecursively(index_arr, slicers, if_decrease_axis);
}

function getSliceIndexRecursively(index_arr, slicers: [Number, Number][], if_decrease_axis: Boolean[]) {
    if (slicers.length === 0) {
        return index_arr; // 如果不再需要切片了就直接返回
    }
    const [start, end] = slicers.shift(); // 当前维度需要切片的开始位置和结束位置
    if (typeof index_arr[0] === 'number') {
        return index_arr.slice(start, end); // 如果只有一个维度就直接切片
    }
    const decrease_next_axis = if_decrease_axis.shift(); // 是否将此维度与下一维度合并
    let sliced_arr = index_arr.slice(start, end).map(sub_arr => getSliceIndexRecursively(
        sub_arr, [...slicers],[...if_decrease_axis])); // 解包后重新组合成为原先的复制而不是引用，slicers每次两个一起pop出来且不会分别修改所以不用深复制
    if (decrease_next_axis) {
        sliced_arr = sliced_arr.flat(); // 合并此维度与下一维度
    }
    return sliced_arr;
}


function mainFunc(
    {
        out, origin
    },
    {
        axes,
        starts,
        ends,
        decrease_axis
    }
) {

    if (
        axes.length > 1
        || starts.length > 1
        || ends.length > 1
        || (decrease_axis && decrease_axis.length === 0)
    ) {
        throw Error('[slice op feature]: current support one dim, support decrease_axis');
    }
    const {
        width_shape,
        height_shape,
        channel,
        total_shape,
        length_unformatted_shape
    } = origin;
    const batch = total_shape / (width_shape * height_shape * channel);
    const tensor_shape = [batch, channel, height_shape, width_shape];

    // let axis = axes[0];

    const new_axes = axes.map(_axis => {
        let axis = _axis;
        if (axis < 0) {
            axis = axis + length_unformatted_shape; // 去掉了+1
        }
        axis = 4 - length_unformatted_shape + axis;
        return axis;
    });
    //
    // if (axis !== 4) {
    //     throw Error('[slice op feature]: unsupport axis value');
    // }
    //
    // const start = starts[0];
    // const end = ends[0];
    //
    // const [
    //     batch_num,
    //     channel_num,
    //     height_num,
    //     width_num
    // ] = tensor_shape;
    //
    // // 计算 output tensor value 对应的 origin index
    // const res_pos = [];
    // for (let index = start; index < end; index++) {
    //     for (let batch = 0; batch < batch_num; batch++) {
    //         for (let channel = 0; channel < channel_num; channel++) {
    //             for (let height = 0; height < height_num; height++) {
    //                 res_pos.push(
    //                     batch * channel_num * height_num * width_num
    //                     + channel * height_num * width_num
    //                     + height * width_num + index
    //                 );
    //             }
    //         }
    //     }
    // }
    const res_pos = getSliceIndex({ axes: new_axes, starts, ends, decrease_axis: [1, 2, 3] }, { shape: tensor_shape }); // 扁平化所有维度
    console.log(tensor_shape, [out.length_unformatted_shape, out.channel, out.height_shape, out.width_shape], new_axes, axes, starts, ends, decrease_axis, res_pos);

    const glslIndexArr = genGLSLArr(res_pos, 'arr', ArrTypeEnum.INT_TYPE);

    const getValueFromArrByIndexGLSL = getValueFromArrByIndex(res_pos, 'arr', ArrTypeEnum.INT_TYPE);

    return `
    ${getValueFromArrByIndexGLSL}

    void main(void) {
        ivec4 oPos = getOutputTensorPos();
        ${glslIndexArr}

        // 输出坐标转换为输入坐标
        int sumVal = oPos.a
            + oPos.b * ${out.width_shape}
            + oPos.g * ${out.height_shape} * ${out.width_shape}
            + oPos.r * ${out.channel} * ${out.width_shape} * ${out.height_shape};

        int index = getValueFromArrByIndex_arr(arr, sumVal);

        float res = 0.0;
        ivec4 co = getTensorPosFromArrayIndex_origin(index);
        res = getValueFromTensorPos_origin(co.r, co.g, co.b, co.a);
        setOutput(float(res));
    }
    `;
}
export default {
    mainFunc,
    textureFuncConf: {
        origin: ['getValueFromTensorPos', 'getTensorPosFromArrayIndex']
    }
};
